package com.imagination.applet.domain.adapter.facechain.impl;

import com.alibaba.fastjson.JSONObject;
import com.imagination.applet.domain.adapter.facechain.FaceChainAdapter;
import com.imagination.applet.domain.aiimage.model.AIImage;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Component
@RequiredArgsConstructor
@Slf4j
public class BaseFaceChainAdapter implements FaceChainAdapter {
    private static final String TASK_TRAIN = "train";
    private static final String TASK_INFERENCE = "inference";

    private final RestTemplate restTemplate;

    @Value(value = "${facechain.api}")
    private String TRAIN_URL;

    @Override
    public void training(Long virtualIdentityId, List<String> images) {
        Map<String, Object> params = new HashMap<>();
        params.put("virtual_identity_id", "virtual-identity-" + virtualIdentityId);
        params.put("instance_image_urls", images);
        Object rsp = restTemplate.postForObject(TRAIN_URL + "/virtual-identity", params, Object.class);
        log.info("[training] params = {}, rsp = {}", params, rsp);
    }

    @Override
    public Integer trainingStatus(Long virtualIdentityId) {
        String url = TRAIN_URL + "/virtual-identity/virtual-identity-" + virtualIdentityId + "/process";
        AIRsp rsp = restTemplate.getForObject(url, AIRsp.class);
        log.info("[trainingStatus] url = {}, rsp = {}", url, rsp);
        return Integer.parseInt(rsp.getCode());
    }

    @Override
    public JSONObject imageHD(AIImage aiImage, String taskId, String imagePath, Integer imagination) {
        Map<String, Object> params = new HashMap<>();
        params.put("task_id", taskId);
        params.put("images", new String[]{imagePath});
        params.put("model_name", aiImage.getModelName());
        params.put("prompt", aiImage.getPrompt());
        params.put("negative_prompt", aiImage.getNegativePrompt());
        params.put("sampler", aiImage.getSampler());
        params.put("inference_steps", aiImage.getInferenceSteps());
        params.put("imagination", new BigDecimal(imagination).divide(new BigDecimal(100),1, RoundingMode.HALF_UP));
        params.put("cfg_scale", aiImage.getCfgScale());
        params.put("num_outputs", 1);
        params.put("width", 1200);
        params.put("height", 1600);
        params.put("is_hd", true);
        log.info("请求参数=>{}", JSONObject.toJSONString(params));
        JSONObject rsp = restTemplate.postForObject(TRAIN_URL + "/virtual-identity-ipa-webui/images", params, JSONObject.class);
        log.info("[training] params = {}, rsp = {}", params, rsp);
        return rsp;
    }

    @Override
    public JSONObject inference(AIImage aiImage, String taskId, Long virtualIdentityId,Integer imagination , String image) {
        if(imagination==null){
            imagination=100;
        }
        Map<String, Object> params = new HashMap<>();
        params.put("task_id", taskId);
        params.put("images", new String[]{image});
        params.put("model_name", aiImage.getModelName());
        params.put("prompt", aiImage.getPrompt());
        params.put("negative_prompt", aiImage.getNegativePrompt());
        params.put("sampler", aiImage.getSampler());
        params.put("denoise", aiImage.getDenoise());
        params.put("instant_weight", aiImage.getInstantWeight());
        params.put("controlnet_weight", aiImage.getControlnetWeight());
        params.put("inference_steps", aiImage.getInferenceSteps());
        params.put("imagination", new BigDecimal(imagination).divide(new BigDecimal(100),1, RoundingMode.HALF_UP));
        params.put("cfg_scale", aiImage.getCfgScale());
        params.put("num_outputs", 2);
        params.put("width", 600);
        params.put("height", 800);
        params.put("is_hd", false);
        log.info("请求参数=>{}", JSONObject.toJSONString(params));
//        JSONObject rsp = restTemplate.postForObject(TRAIN_URL + "/virtual-identity-photomaker/images", params, JSONObject.class);
        JSONObject rsp = restTemplate.postForObject(TRAIN_URL + "/virtual-identity-ipa-webui/images", params, JSONObject.class);
        log.info("[training] params = {}, rsp = {}", params, rsp);
        return rsp;
    }

    @Override
    public JSONObject inferenceStatus(String taskId,String modelName, boolean isHd) {
        String url = TRAIN_URL + "/virtual-identity-ipa/images?task_id=" + taskId+"&model_name="+modelName + "&is_hd=" + isHd;
        JSONObject rsp = restTemplate.getForObject(url, JSONObject.class);
        log.info("[trainingStatus] url = {}, rsp = {}", url, rsp);
        return rsp;
    }

    @Override
    public JSONObject crop(String key, Integer[] arr) {
        Map<String, Object> params = new HashMap<>();
        params.put("file", key);
        params.put("required_shape", arr);
        String url = TRAIN_URL + "/image-crop";
        JSONObject rsp = restTemplate.postForObject(url, params, JSONObject.class);
        log.info("[trainingStatus] url = {}, rsp = {}", url, rsp);
        return rsp;
    }

    @Data
    static class AIRsp {
        private String code;
        private String message;
        private Object data;
    }

}
